import argparse
import glob
import json
import os
import re
from typing import List

import gradio as gr
from llama_index.core.schema import TextNode

from iflytech_assistant.es import index, search
from iflytech_assistant.ugc.avatar import base64_to_image
from iflytech_assistant.ugc.creator import generate_role_single_turn
from iflytech_assistant.ugc.role import Example, RoleMeta
from iflytech_assistant.ugc.role_play import generate

parser = argparse.ArgumentParser()
parser.add_argument("folder", type=str, help="保存微博爬虫爬取结果的文件夹")
parser.add_argument("--do-index", action="store_true", help="是否重建索引")
args = parser.parse_args()


def line_filter(line: str) -> str:
    line = re.sub(r"\s*\[组图.*?\]\s*原图\s*", "", line)
    line = re.sub(r"\s+原图\s+", "", line)
    line = re.sub(r"\s+", " ", line)
    line = line.strip()
    return line


def load_lines_from_folder(folder: str):
    name_2_lines = {}
    for folder_name in os.listdir(folder):
        cur_folder = os.path.join(folder, folder_name)
        if not os.path.isdir(cur_folder):
            continue
        json_files = glob.glob(os.path.join(cur_folder, "*.json"))
        if not json_files:
            continue
        json_file = json_files[0]
        with open(json_file, "r") as f:
            data = json.load(f)
            lines = [line_filter(item["content"]) for item in data["weibo"]]
            lines = [line for line in lines if len(line) > 3]
        name_2_lines[folder_name] = lines
    return name_2_lines


INDEX_NAME = "ugc_weibo_rag"

name_2_lines = load_lines_from_folder(args.folder)


def do_index():
    for name, lines in name_2_lines.items():
        print(f"Indexing {name}")
        nodes = [TextNode(text=line, metadata={"name": name}) for line in lines]
        index(nodes, index_name=INDEX_NAME)


if args.do_index:
    do_index()


def create_role(
    state: dict,
    name: str,
    introduction: str,
    do_web_search: bool,
):
    role = generate_role_single_turn(
        name, do_web_search=do_web_search, generate_avatar=True
    )
    state[name] = role
    return (
        state,
        gr.Dropdown(label="选择角色", choices=state.keys()),
        role.introduction,
        role.category,
        role.appearance,
        role.profile,
        str(role.search_result),
        role.search_time + role.prefilling_time,
        role.generation_time,
        base64_to_image(role.avatar),
    )


def rewrite_message(
    state: dict,
    name: str,
    user_input: str,
):
    role: RoleMeta = state.get(name).model_copy()
    lines: List[TextNode] = search(
        user_input, index_name=INDEX_NAME, rerank=False, filters={"name": name}
    )
    role.lines = [line.text for line in lines]
    example: Example = generate(role, user_input=user_input, _print=False)
    return (
        example.role_output,
        example.system_prompt + "\n" + example.user_prompt,
        role.lines,
    )


def on_selected_role_change(state: dict, name: str):
    role: RoleMeta = state[name]
    return (
        role.name,
        role.introduction,
        role.category,
        role.appearance,
        str(role.search_result),
        base64_to_image(role.avatar),
    )


# Gradio 界面
with gr.Blocks() as demo:
    agent_state = gr.State(value={})

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 角色配置")
            name = gr.Dropdown(label="选择角色", choices=name_2_lines.keys())
            do_web_search = gr.Checkbox(
                label="是否进行网络搜索", value=True, visible=False
            )
            profile = gr.Textbox(label="一句话描述")
            appearance = gr.Textbox(label="外貌特征")
            avatar = gr.Image(label="头像")
            introduction = gr.Textbox(label="角色介绍")
            category = gr.Textbox(label="角色类别", interactive=False)
            search_result = gr.Textbox(label="搜索结果", interactive=False, lines=3)
            create_button = gr.Button("创建角色")
            search_time = gr.Textbox(label="搜索时间", interactive=False)
            create_time = gr.Textbox(label="创建时间", interactive=False)

        with gr.Column():
            gr.Markdown("### 角色扮演")
            role_name = gr.Dropdown(label="选择角色", choices=[])
            lines = gr.Textbox(label="微博台词", lines=5, interactive=False)
            user_input = gr.Textbox(label="输入")
            output = gr.Textbox(label="输出", interactive=False, lines=5)
            rewrite_button = gr.Button("生成")
            raw_prompt = gr.Textbox(label="提示词", interactive=False, lines=10)

    create_button.click(
        fn=create_role,
        inputs=[agent_state, name, introduction, do_web_search],
        outputs=[
            agent_state,
            role_name,
            introduction,
            category,
            appearance,
            profile,
            search_result,
            search_time,
            create_time,
            avatar,
        ],
    )

    # on role_name change
    role_name.select(
        fn=on_selected_role_change,
        inputs=[agent_state, role_name],
        outputs=[name, introduction, category, appearance, search_result, avatar],
    )

    rewrite_button.click(
        fn=rewrite_message,
        inputs=[agent_state, role_name, user_input],
        outputs=[output, raw_prompt, lines],
    )


if __name__ == "__main__":
    demo.launch(share=True)
